#!/usr/bin/env python3
"""
腾讯云COS上传脚本
将HLS文件上传到COS并配置CDN
"""

import os
import sys
import json
import argparse
from pathlib import Path
from qcloud_cos import CosConfig
from qcloud_cos import CosS3Client
from qcloud_cos.cos_exception import CosServiceError, CosClientError

def load_config(config_file='config.json'):
    """加载配置文件"""
    try:
        script_dir = os.path.dirname(os.path.abspath(__file__))
        config_path = os.path.join(script_dir, config_file)

        if not os.path.exists(config_path):
            return None

        with open(config_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    except Exception as e:
        print(f"读取配置文件失败: {e}")
        return None

def get_next_number(base_prefix, client, bucket):
    """获取下一个可用的文件夹编号"""
    if not base_prefix.endswith('_auto'):
        return None, base_prefix

    # 如果以_auto结尾，则自动生成递增的前缀
    base_name = base_prefix.replace('_auto', '')
    existing_numbers = []

    try:
        # 检查COS上现有的对象
        response = client.list_objects(
            Bucket=bucket,
            Prefix=f"{base_name}_",
            Delimiter='/'
        )

        if 'CommonPrefixes' in response:
            for prefix_info in response['CommonPrefixes']:
                prefix = prefix_info['Prefix']
                if prefix.startswith(f"{base_name}_"):
                    suffix = prefix[len(f"{base_name}_"):].rstrip('/')
                    if suffix.isdigit():
                        existing_numbers.append(int(suffix))

        # 检查本地目录
        local_base = os.path.dirname(base_name)
        if os.path.exists(local_base):
            for name in os.listdir(local_base):
                if name.startswith(os.path.basename(base_name) + "_"):
                    suffix = name[len(os.path.basename(base_name)) + 1:]
                    if suffix.isdigit():
                        existing_numbers.append(int(suffix))

    except Exception as e:
        print(f"❌ 检查现有前缀时出错: {e}")
        existing_numbers = []

    # 确定下一个编号
    next_number = max(existing_numbers or [1]) + 1
    final_prefix = f"{base_name}_{next_number}"

    print(f"🔍 检测到现有文件夹: {sorted(existing_numbers) if existing_numbers else '无'}")
    print(f"📈 下一个编号: {next_number}")
    print(f"🔄 自动递增: {base_prefix} -> {final_prefix}")

    return next_number, final_prefix

class COSUploader:
    def __init__(self):
        # 加载配置
        self.config = load_config()
        if not self.config:
            raise ValueError("无法加载配置文件")

        cos_config = self.config.get('cos', {})
        self.secret_id = cos_config.get('secret_id')
        self.secret_key = cos_config.get('secret_key')
        self.region = cos_config.get('region')
        self.bucket = cos_config.get('bucket')
        self.auto_increment = cos_config.get('auto_increment', True)
        self.retry_times = cos_config.get('retry_times', 3)
        self.part_size = cos_config.get('part_size', 1)
        self.max_thread = cos_config.get('max_thread', 10)

        cdn_config = self.config.get('cdn', {})
        self.cdn_domain = cdn_config.get('domain')
        self.cdn_protocol = cdn_config.get('protocol', 'https')
        self.cdn_enabled = cdn_config.get('enable', True)

        if not all([self.secret_id, self.secret_key, self.region, self.bucket]):
            raise ValueError("COS配置不完整")

        # 初始化COS客户端
        config = CosConfig(
            Region=self.region,
            SecretId=self.secret_id,
            SecretKey=self.secret_key,
            Token=None,
            Scheme=self.cdn_protocol
        )
        self.client = CosS3Client(config)
        
    def upload_file(self, local_path, cos_key):
        """上传单个文件到COS"""
        for retry in range(self.retry_times):
            try:
                self.client.upload_file(
                    Bucket=self.bucket,
                    LocalFilePath=local_path,
                    Key=cos_key,
                    PartSize=self.part_size,
                    MAXThread=self.max_thread
                )
                print(f"✅ 上传成功: {local_path} -> {cos_key}")
                return True
            except (CosServiceError, CosClientError) as e:
                print(f"❌ 上传失败 {local_path} (重试 {retry + 1}/{self.retry_times}): {e}")
                if retry == self.retry_times - 1:
                    return False
                continue

    def upload_hls_directory(self, local_dir, cos_prefix):
        """上传整个HLS目录"""
        local_dir = Path(local_dir)
        success_count = 0
        total_count = 0
        
        # 上传所有文件
        for file_path in local_dir.rglob('*'):
            if file_path.is_file():
                relative_path = file_path.relative_to(local_dir)
                cos_key = f"{cos_prefix}/{relative_path}".replace('\\', '/')
                
                total_count += 1
                if self.upload_file(str(file_path), cos_key):
                    success_count += 1
        
        print(f"\n📊 上传完成: {success_count}/{total_count} 文件成功")
        return success_count == total_count
    
    def set_cors_policy(self):
        """设置CORS策略以支持HLS播放"""
        cors_config = {
            'CORSRule': [
                {
                    'ID': 'HLS-CORS',
                    'AllowedOrigin': ['*'],
                    'AllowedMethod': ['GET', 'HEAD'],
                    'AllowedHeader': ['*'],
                    'MaxAgeSeconds': 3600
                }
            ]
        }
        
        try:
            self.client.put_bucket_cors(
                Bucket=self.bucket,
                CORSConfiguration=cors_config
            )
            print("✅ CORS策略设置成功")
            return True
        except Exception as e:
            print(f"❌ CORS策略设置失败: {e}")
            return False
    
    def generate_cdn_urls(self, cos_prefix):
        """生成CDN访问URL"""
        if self.cdn_enabled and self.cdn_domain:
            base_url = f"{self.cdn_protocol}://{self.cdn_domain}/{cos_prefix}"
        else:
            base_url = f"{self.cdn_protocol}://{self.bucket}.cos.{self.region}.myqcloud.com/{cos_prefix}"

        urls = {
            "playlist_url": f"{base_url}/stream.m3u8",
            "key_url": f"{base_url}/encryption.key",
            "base_url": base_url
        }
        
        return urls

    def get_next_cos_prefix(self, base_prefix):
        """获取下一个可用的COS前缀"""
        if not self.auto_increment or not base_prefix.endswith('_auto'):
            return base_prefix

        # 如果以_auto结尾，则自动生成递增的前缀
        base_name = base_prefix.replace('_auto', '')
        existing_numbers = []

        try:
            # 1. 检查COS上已有的对象
            response = self.client.list_objects(
                Bucket=self.bucket,
                Prefix=f"{base_name}_",
                Delimiter='/'
            )

            if 'CommonPrefixes' in response:
                for prefix_info in response['CommonPrefixes']:
                    prefix = prefix_info['Prefix']
                    if prefix.startswith(f"{base_name}_"):
                        suffix = prefix[len(f"{base_name}_"):].rstrip('/')
                        if suffix.isdigit():
                            existing_numbers.append(int(suffix))

            print(f"COS中检测到的编号: {sorted(existing_numbers)}")

            # 2. 检查本地目录中的编号
            local_base_dir = os.path.dirname(base_name)
            if os.path.exists(local_base_dir):
                local_base_name = os.path.basename(base_name)
                for name in os.listdir(local_base_dir):
                    if name.startswith(f"{local_base_name}_"):
                        suffix = name[len(local_base_name) + 1:]
                        if suffix.isdigit():
                            existing_numbers.append(int(suffix))

            print(f"本地检测到的编号: {sorted(existing_numbers)}")

        except Exception as e:
            print(f"❌ 检查现有前缀时出错: {e}")

        # 3. 确定下一个编号
        next_number = max(existing_numbers or [1]) + 1
        final_prefix = f"{base_name}_{next_number}"

        print(f"🔍 检测到现有编号: {sorted(existing_numbers) if existing_numbers else '无'}")
        print(f"📈 下一个编号: {next_number}")
        print(f"🔄 自动递增: {base_prefix} -> {final_prefix}")

        return final_prefix

def main():
    parser = argparse.ArgumentParser(description='COS上传工具')
    parser.add_argument('local_dir', help='本地HLS目录路径')
    parser.add_argument('cos_prefix', help='COS存储路径前缀 (使用_auto后缀可自动递增)')
    parser.add_argument('--no-cdn', action='store_true', help='不使用CDN')
    parser.add_argument('--dry-run', action='store_true', help='仅预览不上传')

    args = parser.parse_args()

    try:
        uploader = COSUploader()
    except Exception as e:
        print(f"❌ 初始化失败: {e}")
        sys.exit(1)
    
    # 检查本地目录
    if not os.path.exists(args.local_dir):
        print(f"❌ 本地目录不存在: {args.local_dir}")
        sys.exit(1)
    
    # 获取实际的COS前缀
    final_cos_prefix = uploader.get_next_cos_prefix(args.cos_prefix)

    print(f"\n📤 准备上传:")
    print(f"   本地目录: {args.local_dir}")
    print(f"   COS前缀: {final_cos_prefix}")

    if args.dry_run:
        print("\n🔍 预览模式 - 将要上传的文件:")
        local_dir = Path(args.local_dir)
        for file_path in local_dir.rglob('*'):
            if file_path.is_file():
                relative_path = file_path.relative_to(local_dir)
                cos_key = f"{final_cos_prefix}/{relative_path}".replace('\\', '/')
                print(f"   {file_path.name} -> {cos_key}")
        return

    # 设置CORS策略
    uploader.set_cors_policy()

    # 上传文件
    if uploader.upload_hls_directory(args.local_dir, final_cos_prefix):
        # 生成访问URL
        urls = uploader.generate_cdn_urls(final_cos_prefix)

        print("\n=== 上传完成 ===")
        print(f"播放列表URL: {urls['playlist_url']}")
        print(f"密钥URL: {urls['key_url']}")
        print(f"基础URL: {urls['base_url']}")
        
        # 保存URL信息
        url_file = Path(args.local_dir) / 'urls.json'
        with open(url_file, 'w', encoding='utf-8') as f:
            json.dump(urls, f, indent=2, ensure_ascii=False)
        print(f"\n✅ URL信息已保存到: {url_file}")
    else:
        print("\n❌ 上传失败!")
        sys.exit(1)

if __name__ == "__main__":
    main()
